#!/usr/bin/env python
# -*- encoding: utf-8 -*-
# Project: spd-sxmcc
"""
@author: lyndon
@time Created on 2018/10/9 17:46
@desc
"""

from es.errors import *
from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk, reindex
import json
import time
import argparse
import sys

reload(sys)
sys.setdefaultencoding('utf-8')

# ES索引和Type名称
INDEX_NAME = "twitter"
TYPE_NAME = "tweet"

my_settings = {
    "settings": {
        "number_of_shards": "5",
        "number_of_replicas": "1",
    },
}


class EsTool:
    """# ES操作工具类"""

    # 类初始化函数
    def __init__(self, hosts, timeout):
        self.es = Elasticsearch(hosts, timeout=timeout)

    def set_from_dict(self, dict_data=None, index_name=None, doc_type_name=None):
        """将字典存储到es中"""
        if index_name is None:
            raise NotHasIndexError('index_name not specified')

        if doc_type_name is None:
            raise NotHasTypeError('doc_type_name not specified')

        if dict_data is None:
            dict_data = {}
        mp_id = dict_data.pop("_id")
        body = json.dumps(dict_data, ensure_ascii=False)
        ret = self.es.index(index=index_name, doc_type=doc_type_name, id=mp_id, body=body)
        return dict(ret).get('result')

    def set_from_dict_array(self, row_dict_data, index_name=INDEX_NAME, doc_type_name=TYPE_NAME):
        """将数组中的dict存储到es中"""

        if index_name is None:
            raise NotHasIndexError('index_name not specified')

        if doc_type_name is None:
            raise NotHasTypeError('doc_type_name not specified')

        if row_dict_data is None:
            row_dict_data = []

        # 创建ACTIONS
        ACTIONS = []
        # print "es set_data length",len(fields_data)
        for dict_data in row_dict_data:
            assert type(dict_data).__name__ == 'dict'
            # print "fields", fields
            # print fields[1]
            action = {
                "_index": index_name,
                "_type": doc_type_name,
                "_source": dict_data
            }
            ACTIONS.append(action)

        # print "len ACTIONS", len(ACTIONS)
        # 批量处理
        # success, _ = bulk(self.es, ACTIONS, index=index_name, raise_on_error=True)
        # print('Performed %d actions' % success)
        ret = bulk(self.es, ACTIONS, index=index_name, raise_on_error=True)
        return ret

    # 读取参数
    @classmethod
    def read_args(cls):
        parser = argparse.ArgumentParser(description="Search Elastic Engine")
        parser.add_argument("-i", dest="input_file", action="store", help="input file1", required=False,
                            default="./data.txt")
        # parser.add_argument("-o", dest="output_file", action="store", help="output file", required=True)
        return parser.parse_args()

    @staticmethod
    def select_mapping(index_name):
        global my_mapping
        if index_name == 'tb_phone_url':
            my_mapping = {
                "properties": {
                    "article_id": {
                        "type": "keyword"
                    },
                    "msisdn": {
                        "type": "text",
                        "analyzer": "ik_max_word"
                    },
                    "url": {
                        "type": "text",
                        "analyzer": "ik_max_word"
                    },
                    "deal_date": {
                        "type": "text"
                    }
                }
            }
        elif index_name == 'tb_library':
            my_mapping = {
                "properties": {
                    "title": {
                        "type": "text",
                        "analyzer": "ik_max_word"
                    },
                    "weixin_mp_name": {
                        "type": "text",
                        "analyzer": "ik_max_word"
                    },
                    "content": {
                        "type": "text",
                        "analyzer": "ik_max_word"
                    },
                    "extract_time": {
                        "type": "date"
                    }
                }
            }
        return my_mapping

    # 初始化index，设置mapping
    def init_es(self, index_name=INDEX_NAME, doc_type_name=TYPE_NAME, mapping=None):

        if mapping is None:
            mapping = self.select_mapping(index_name)
        try:
            # 先销毁，后创建Index和mapping
            try:
                delete_index = self.es.indices.delete(index=index_name)  # {u'acknowledged': True}
            except Exception as delete_e:
                print('delete exception:', delete_e)
                pass
            # create_index = self.es.indices.create(index=index_name, body=my_mapping)  # {u'acknowledged': True}
            create_index = self.es.indices.create(index=index_name, body=my_settings)
            print(create_index)
            # if delete_index["acknowledged"] != True or create_index["acknowledged"] != True:
            #     print("Index creation failed...")
            mapping_index = self.es.indices.put_mapping(doc_type=doc_type_name,
                                                        body=mapping, index=index_name)  # {u'acknowledged': True}
            if create_index["acknowledged"] != True or mapping_index["acknowledged"] != True:
                print("Index creation failed...")
        except Exception as e:
            print("set_mapping except", e)

    def get_search(self, index_name=INDEX_NAME, doc_type_name=TYPE_NAME, body=None, params=None):
        """
        Execute a search query and get back search hits that match the query.
        `<http://www.elastic.co/guide/en/elasticsearch/reference/current/search-search.html>`_

        :arg index: A comma-separated list of index names to search; use `_all`
            or empty string to perform the operation on all indices
        :arg doc_type: A comma-separated list of document types to search; leave
            empty to perform the operation on all types
        :arg body: The search definition using the Query DSL
        :arg _source: True or false to return the _source field or not, or a
            list of fields to return
        :arg _source_exclude: A list of fields to exclude from the returned
            _source field
        :arg _source_include: A list of fields to extract and return from the
            _source field
        :arg allow_no_indices: Whether to ignore if a wildcard indices
            expression resolves into no concrete indices. (This includes `_all`
            string or when no indices have been specified)
        :arg allow_partial_search_results: Set to false to return an overall
            failure if the request would produce partial results. Defaults to
            True, which will allow partial results in the case of timeouts or
            partial failures
        :arg analyze_wildcard: Specify whether wildcard and prefix queries
            should be analyzed (default: false)
        :arg analyzer: The analyzer to use for the query string
        :arg batched_reduce_size: The number of shard results that should be
            reduced at once on the coordinating node. This value should be used
            as a protection mechanism to reduce the memory overhead per search
            request if the potential number of shards in the request can be
            large., default 512
        :arg default_operator: The default operator for query string query (AND
            or OR), default 'OR', valid choices are: 'AND', 'OR'
        :arg df: The field to use as default where no field prefix is given in
            the query string
        :arg docvalue_fields: A comma-separated list of fields to return as the
            docvalue representation of a field for each hit
        :arg expand_wildcards: Whether to expand wildcard expression to concrete
            indices that are open, closed or both., default 'open', valid
            choices are: 'open', 'closed', 'none', 'all'
        :arg explain: Specify whether to return detailed information about score
            computation as part of a hit
        :arg from\\_: Starting offset (default: 0)
        :arg ignore_unavailable: Whether specified concrete indices should be
            ignored when unavailable (missing or closed)
        :arg lenient: Specify whether format-based query failures (such as
            providing text to a numeric field) should be ignored
        :arg max_concurrent_shard_requests: The number of concurrent shard
            requests this search executes concurrently. This value should be
            used to limit the impact of the search on the cluster in order to
            limit the number of concurrent shard requests, default 'The default
            grows with the number of nodes in the cluster but is at most 256.'
        :arg pre_filter_shard_size: A threshold that enforces a pre-filter
            roundtrip to prefilter search shards based on query rewriting if
            the number of shards the search request expands to exceeds the
            threshold. This filter roundtrip can limit the number of shards
            significantly if for instance a shard can not match any documents
            based on it's rewrite method ie. if date filters are mandatory to
            match but the shard bounds and the query are disjoint., default 128
        :arg preference: Specify the node or shard the operation should be
            performed on (default: random)
        :arg q: Query in the Lucene query string syntax
        :arg request_cache: Specify if request cache should be used for this
            request or not, defaults to index level setting
        :arg routing: A comma-separated list of specific routing values
        :arg scroll: Specify how long a consistent view of the index should be
            maintained for scrolled search
        :arg search_type: Search operation type, valid choices are:
            'query_then_fetch', 'dfs_query_then_fetch'
        :arg size: Number of hits to return (default: 10)
        :arg sort: A comma-separated list of <field>:<direction> pairs
        :arg stats: Specific 'tag' of the request for logging and statistical
            purposes
        :arg stored_fields: A comma-separated list of stored fields to return as
            part of a hit
        :arg suggest_field: Specify which field to use for suggestions
        :arg suggest_mode: Specify suggest mode, default 'missing', valid
            choices are: 'missing', 'popular', 'always'
        :arg suggest_size: How many suggestions to return in response
        :arg suggest_text: The source text for which the suggestions should be
            returned
        :arg terminate_after: The maximum number of documents to collect for
            each shard, upon reaching which the query execution will terminate
            early.
        :arg timeout: Explicit operation timeout
        :arg track_scores: Whether to calculate and return scores even if they
            are not used for sorting
        :arg track_total_hits: Indicate if the number of documents that match
            the query should be tracked
        :arg typed_keys: Specify whether aggregation and suggester names should
            be prefixed by their respective types in the response
        :arg version: Specify whether to return document version as part of a
            hit
        """

        if params is None:
            return self.es.search(index=index_name, doc_type=doc_type_name, body=body)
        else:
            return self.es.search(index=index_name, doc_type=doc_type_name, body=body, params=params)

    def collect_highlight(self, ret_dict):
        print(ret_dict)
        try:
            ret_dict = dict(ret_dict)
        except Exception as e:
            print('return value not a dict', e)
            pass

        # print(ret_dict[])
        print('---------------------------------')
        print(ret_dict['hits']['hits'])
        try:
            if ret_dict['timed_out'] == False and ret_dict['hits'] is not None:
                if ret_dict['hits']['hits'] is not None:
                    highlights = ret_dict['hits']['hits']
                    for hgls in highlights:
                        content = hgls['highlight']['content']
                        yield content
        except Exception as e:
            raise e

    # 主函数
if __name__ == '__main__':
    # args = read_args()
    # 初始化es环境
    # init_es(hosts=["192.168.100.181","192.168.100.182","192.168.100.183"], timeout=5000)



    # 执行写入操作
    # tweet_list = [("111","222","333","444","555"), ("11","22","33","44","55")]
    # es.set_data(tweet_list)

    # 创建es类
    es = EsTool(hosts=["192.168.20.110", "192.168.20.111", "192.168.20.112"], timeout=5000)
    # 初始化index
    es.init_es(index_name='tb_library', doc_type_name='_doc')
    es.init_es(index_name='tb_phone_url', doc_type_name='_doc')

    # 设定mapping

    # 创建index
    # es.init_es(index_name='tb_index', doc_type_name='_doc')

    # data = {'article_id': 'dafarr', 'msisdn': '13648917112', 'content_name': '大口大口去', 'money': 1918}
    # data = None
    # ret = es.set_from_dict(dict_data=data, index_name='tb_index', doc_type_name='_doc')
    # print(ret)

    # data = [{'article_id': 'dfaf', 'msisdn': '19181717111', 'content_name': '发发健康', 'money': 19181},
    #         {'article_id': 'daed3', 'msisdn': '13119181881', 'content_name': '搭客气话', 'money': 10189}, ]
    # ret = es.set_from_dict_array(data, index_name='tb_index', doc_type_name='_doc')
    # print(ret)
    # es.init_es('tb_library', '_doc')
